"""
核反应堆子通道分析程序的文件操作函数
"""

import os
import csv
import numpy as np
import json
import traceback

def create_results_folder(folder_name):
    """
    创建结果保存文件夹
    
    参数:
    folder_name: str, 文件夹名称
    """
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)

def save_to_csv(analysis, time_step):
    """
    保存当前时间步的数据到CSV文件
    
    参数:
    analysis: SubchannelAnalysis对象
    time_step: int, 当前时间步编号
    """
    try:
        filename = os.path.join(analysis.results_folder, f"timestep_{time_step:06d}.csv")
        
        # 使用UTF-8编码保存
        with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            
            # 1. 写入时间步信息
            writer.writerow(['时间步信息'])
            writer.writerow(['当前时间 (s)', f"{analysis.current_time:.6f}"])
            writer.writerow(['时间步编号', time_step])
            writer.writerow([])
            
            # 2. 写入全局参数
            writer.writerow(['全局参数'])
            writer.writerow(['平均压力 (Pa)', f"{np.mean(analysis.pressure):.2f}"])
            writer.writerow(['最大温度 (K)', f"{np.max(analysis.temperature):.2f}"])
            writer.writerow(['最大空泡率', f"{np.max(analysis.void_fraction):.4f}"])
            writer.writerow([])
            
            # 3. 写入详细数据
            headers = [
                '轴向位置 (m)',
                '子通道编号',
                '空泡率',
                '压力 (Pa)',
                '液相温度 (K)',
                '气相温度 (K)',
                '壁面温度 (K)',
                '液相速度 (m/s)',
                '气相速度 (m/s)',
                '液相壁面热流密度 (W/m²)',
                '气相壁面热流密度 (W/m²)',
                '相变率 (kg/m³·s)'
            ]
            writer.writerow(headers)
            
            # 写入每个网格点的数据
            for i in range(analysis.n_axial):
                for j in range(analysis.n_channels):
                    row_data = [
                        f"{i * analysis.dz:.6f}",
                        j + 1,
                        f"{analysis.void_fraction[i, j]:.6f}",
                        f"{analysis.pressure[i, j]:.2f}",
                        f"{analysis.T_l[i, j]:.2f}",
                        f"{analysis.T_v[i, j]:.2f}",
                        f"{analysis.T_c[i, j]:.2f}",
                        f"{analysis.liquid_velocity[i, j]:.4f}",
                        f"{analysis.vapor_velocity[i, j]:.4f}",
                        f"{analysis.q_wl[i, j]:.2f}",
                        f"{analysis.q_wv[i, j]:.2f}",
                        f"{analysis.gamma[i, j]:.6f}"
                    ]
                    writer.writerow(row_data)
                    
        print(f"结果已保存到: {filename}")
        
    except Exception as e:
        print(f"保存CSV文件时出错: {str(e)}")

def save_full_state(analysis, time_step):
    """
    保存完整的计算状态
    
    参数:
    analysis: SubchannelAnalysis对象
    time_step: int, 当前时间步编号
    """
    try:
        filename = os.path.join(analysis.results_folder, f"full_state_{time_step:06d}.npz")
        np.savez_compressed(
            filename,
            time=analysis.current_time,
            time_step=time_step,
            void_fraction=analysis.void_fraction,
            pressure=analysis.pressure,
            temperature=analysis.temperature,
            liquid_velocity=analysis.liquid_velocity,
            vapor_velocity=analysis.vapor_velocity,
            T_l=analysis.T_l,
            T_v=analysis.T_v,
            T_c=analysis.T_c,
            q_wl=analysis.q_wl,
            q_wv=analysis.q_wv,
            gamma=analysis.gamma
        )
        print(f"完整状态已保存到: {filename}")
        
    except Exception as e:
        print(f"保存完整状态时出错: {str(e)}")

def save_history_data(analysis):
    """
    保存历史数据到JSON文件
    
    参数:
    analysis: SubchannelAnalysis对象
    """
    try:
        filename = os.path.join(analysis.results_folder, "history.json")
        
        # 将NumPy数组转换为列表
        history_data = {}
        for key, value in analysis.history.items():
            if isinstance(value, list):
                if isinstance(value[0], np.ndarray):
                    history_data[key] = [arr.tolist() for arr in value]
                else:
                    history_data[key] = value
            elif isinstance(value, np.ndarray):
                history_data[key] = value.tolist()
            else:
                history_data[key] = value
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(history_data, f, indent=4, ensure_ascii=False)
            
        print(f"历史数据已保存到: {filename}")
        
    except Exception as e:
        print(f"保存历史数据时出错: {str(e)}")

def load_full_state(filename):
    """
    加载完整的计算状态
    
    参数:
    filename: str, 状态文件路径
    
    返回:
    dict: 包含所有状态变量的字典
    """
    try:
        data = np.load(filename)
        state = {key: data[key] for key in data.files}
        return state
        
    except Exception as e:
        print(f"加载状态文件时出错: {str(e)}")
        return None

def load_history_data(filename):
    """
    加载历史数据
    
    参数:
    filename: str, 历史数据文件路径
    
    返回:
    dict: 包含所有历史数据的字典
    """
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            history_data = json.load(f)
            
        # 将列表转换回NumPy数组
        for key, value in history_data.items():
            if isinstance(value, list):
                if isinstance(value[0], list):
                    history_data[key] = [np.array(arr) for arr in value]
                else:
                    history_data[key] = np.array(value)
                    
        return history_data
        
    except Exception as e:
        print(f"加载历史数据时出错: {str(e)}")
        return None

def export_results_summary(analysis, filename):
    """
    导出计算结果摘要
    
    参数:
    analysis: SubchannelAnalysis对象
    filename: str, 输出文件名
    """
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            # 写入基本信息
            f.write("核反应堆子通道分析结果摘要\n")
            f.write("=" * 40 + "\n\n")
            
            # 写入计算参数
            f.write("计算参数:\n")
            f.write(f"轴向节点数: {analysis.n_axial}\n")
            f.write(f"子通道数量: {analysis.n_channels}\n")
            f.write(f"通道长度: {analysis.channel_length} m\n")
            f.write(f"时间步长: {analysis.time_step} s\n")
            f.write(f"总计算时间: {analysis.current_time} s\n\n")
            
            # 写入入口条件
            f.write("入口条件:\n")
            f.write(f"入口压力: {analysis.inlet_pressure/1e6:.2f} MPa\n")
            f.write(f"入口温度: {analysis.inlet_temp-273.15:.2f} °C\n")
            f.write(f"入口质量流量: {analysis.inlet_mass_flux:.2f} kg/m²·s\n\n")
            
            # 写入计算结果统计
            f.write("计算结果统计:\n")
            f.write(f"最大空泡率: {np.max(analysis.void_fraction):.4f}\n")
            f.write(f"最大压力: {np.max(analysis.pressure)/1e6:.2f} MPa\n")
            f.write(f"最大温度: {np.max(analysis.temperature)-273.15:.2f} °C\n")
            f.write(f"最大液相速度: {np.max(analysis.liquid_velocity):.2f} m/s\n")
            f.write(f"最大气相速度: {np.max(analysis.vapor_velocity):.2f} m/s\n")
            
        print(f"结果摘要已保存到: {filename}")
        
    except Exception as e:
        print(f"导出结果摘要时出错: {str(e)}")
        print(f"错误位置: {traceback.format_exc()}")

def check_file_integrity(filename):
    """
    检查文件完整性
    
    参数:
    filename: str, 文件路径
    
    返回:
    bool: 文件是否完整
    """
    try:
        if filename.endswith('.npz'):
            # 检查NPZ文件
            data = np.load(filename)
            required_keys = ['time', 'time_step', 'void_fraction', 'pressure', 'temperature']
            return all(key in data.files for key in required_keys)
            
        elif filename.endswith('.json'):
            # 检查JSON文件
            with open(filename, 'r', encoding='utf-8') as f:
                data = json.load(f)
            required_keys = ['time', 'time_steps', 'residuals']
            return all(key in data for key in required_keys)
            
        elif filename.endswith('.csv'):
            # 检查CSV文件
            with open(filename, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                headers = next(reader)
                return len(headers) >= 10  # 确保至少有10列数据
                
        return False
        
    except Exception as e:
        print(f"检查文件完整性时出错: {str(e)}")
        return False
